# -*- coding: utf-8 -*-
import numpy as np
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA

class ClusterFeatures(object):
    """
    Basic handling of clustering features.
    """
    def __init__(self, pca_k, random_state):
        """
        :param pca_k: If you want the features to be ran through pca, this is the components number.
        :param random_state: Random state.
        """
        self.pca_k = pca_k
        self.random_state = random_state

    def __get_model(self, k):
        """
        Retrieve clustering model.
        :param k: amount of clusters.
        :return: Clustering model.
        """
        return KMeans(n_clusters=k, random_state=self.random_state)

    def __get_centroids(self, model):
        """
        Retrieve centroids of model.
        :param model: Clustering model.
        :return: Centroids.
        """
        return model.cluster_centers_

    def __find_closest_args(self, features, centroids):
        """
        Find the closest arguments to centroid.
        :param features: the embedding matrix created by bert parent.
        :param centroids: Centroids to find closest.
        :return: Closest arguments.
        """
        centroid_min = 1e10
        cur_arg = -1
        args = {}
        used_idx = []

        for j, centroid in enumerate(centroids):
            for i, feature in enumerate(features):
                value = np.linalg.norm(feature - centroid)
                if value < centroid_min and i not in used_idx:
                    cur_arg = i
                    centroid_min = value
            used_idx.append(cur_arg)
            args[j] = cur_arg
            centroid_min = 1e10
            cur_arg = -1
        return args

    def cluster(self, features, num_sentences, ratio):
        """
        Clusters sentences based on the ratio.
        :param features: the embedding matrix created by bert parent.
        :param num_sentences: Number of sentences. Overrides ratio.
        :param ratio: Ratio to use for clustering.
        :return: Sentences index that qualify for summary.
        """
        if self.pca_k:
            features = PCA(n_components=self.pca_k).fit_transform(features)
        
        if num_sentences is not None:
            if num_sentences == 0:
                return []
            k = min(num_sentences, len(features))
        else:
            k = max(int(len(features) * ratio), 1)

        model = self.__get_model(k).fit(features)
        centroids = self.__get_centroids(model)
        cluster_args = self.__find_closest_args(features, centroids)
        sorted_values = sorted(cluster_args.values())
        return sorted_values
    
    def __call__(self, features, num_sentences, ratio):
        return self.cluster(features, num_sentences, ratio)
    